# -*- coding:utf-8 -*-
# --------------------------------------------------------
# Fast R-CNN
# Copyright (c) 2015 Microsoft
# Licensed under The MIT License [see LICENSE for details]
# Written by Ross Girshick
# --------------------------------------------------------

"""Test a Fast R-CNN network on an imdb (image database)."""

# -----Newly Added----------------
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
import shutil, time
from collections import OrderedDict
from pandas import DataFrame
# -----------------------------------
from fast_rcnn.config import cfg, get_output_dir
from fast_rcnn.bbox_transform import clip_boxes, bbox_transform_inv
import argparse
from utils.timer import Timer
import numpy as np
import cv2
import caffe
from fast_rcnn.nms_wrapper import nms
import cPickle
from utils.blob import im_list_to_blob
import os
from utils.cython_bbox import bbox_overlaps

# 用来保存评价前的预测框的信息
dict_list_before = []

def _get_image_blob(im):
    """Converts an image into a network input.
    生成对应target_size的图像,本质是调整图像大小且生成符合blob格式的图像数据,即cv2.resize,
    使得图像高与宽其中一个是固定的, 比如是640(cfg.TEST.SCALES)或者2000(cfg.TEST.MAX_SIZE)

    Arguments:
        im (ndarray): a color image in BGR order

    Returns:
        blob (ndarray): a data blob holding an image pyramid
        im_scale_factors (list): list of image scales (relative to im) used
            in the image pyramid
    """
    # ndarray.astype() --> Copy of the array, cast to a specified type.
    im_orig = im.astype(np.float32, copy=True)
    im_orig -= cfg.PIXEL_MEANS

    im_shape = im_orig.shape
    im_size_min = np.min(im_shape[0:2])
    im_size_max = np.max(im_shape[0:2])

    processed_ims = []
    im_scale_factors = []

    for target_size in cfg.TEST.SCALES:
        im_scale = float(target_size) / float(im_size_min)
        # Prevent the biggest axis from being more than MAX_SIZE
        if np.round(im_scale * im_size_max) > cfg.TEST.MAX_SIZE:
            im_scale = float(cfg.TEST.MAX_SIZE) / float(im_size_max)

        # Make width and height be multiples of a specified number
        im_scale_x = np.floor(im.shape[1] * im_scale / cfg.TEST.SCALE_MULTIPLE_OF) * cfg.TEST.SCALE_MULTIPLE_OF / im.shape[1]
        im_scale_y = np.floor(im.shape[0] * im_scale / cfg.TEST.SCALE_MULTIPLE_OF) * cfg.TEST.SCALE_MULTIPLE_OF / im.shape[0]
        im = cv2.resize(im_orig, None, None, fx=im_scale_x, fy=im_scale_y,
                        interpolation=cv2.INTER_LINEAR)
        # print('----im_orig_size is {}, and im_resize_size is {}'.format(im_orig.shape[0:2], im.shape[0:2]))
        im_scale_factors.append(np.array([im_scale_x, im_scale_y, im_scale_x, im_scale_y]))
        processed_ims.append(im)

    # Create a blob to hold the input images
    blob = im_list_to_blob(processed_ims)
    # blob是一个ndarray, 格式为(batch elem, channel, height, width)
    return blob, np.array(im_scale_factors)

def _get_rois_blob(im_rois, im_scale_factors):
    """Converts RoIs into network inputs.

    Arguments:
        im_rois (ndarray): R x 4 matrix of RoIs in original image coordinates
        im_scale_factors (list): scale factors as returned by _get_image_blob

    Returns:
        blob (ndarray): R x 5 matrix of RoIs in the image pyramid
    """
    rois, levels = _project_im_rois(im_rois, im_scale_factors)
    rois_blob = np.hstack((levels, rois))
    return rois_blob.astype(np.float32, copy=False)

def _project_im_rois(im_rois, scales):
    """Project image RoIs into the image pyramid built by _get_image_blob.

    Arguments:
        im_rois (ndarray): R x 4 matrix of RoIs in original image coordinates
        scales (list): scale factors as returned by _get_image_blob

    Returns:
        rois (ndarray): R x 4 matrix of projected RoI coordinates
        levels (list): image pyramid levels used by each projected RoI
    """
    im_rois = im_rois.astype(np.float, copy=False)

    if len(scales) > 1:
        widths = im_rois[:, 2] - im_rois[:, 0] + 1
        heights = im_rois[:, 3] - im_rois[:, 1] + 1

        areas = widths * heights
        scaled_areas = areas[:, np.newaxis] * (scales[np.newaxis, :] ** 2)
        diff_areas = np.abs(scaled_areas - 224 * 224)
        levels = diff_areas.argmin(axis=1)[:, np.newaxis]
    else:
        levels = np.zeros((im_rois.shape[0], 1), dtype=np.int)

    rois = im_rois * scales[levels]

    return rois, levels

def _get_blobs(im, rois):
    """Convert an image and RoIs within that image into network inputs."""
    blobs = {'data': None, 'rois': None}
    blobs['data'], im_scale_factors = _get_image_blob(im)
    if not cfg.TEST.HAS_RPN:
        blobs['rois'] = _get_rois_blob(rois, im_scale_factors)
    return blobs, im_scale_factors

def im_detect(net, im, _t=None, boxes=None):
    """Detect object classes in an image given object proposals.

    Arguments:
        net (caffe.Net): Fast R-CNN network to use
        im (ndarray): color image to test (in BGR order)
        boxes (ndarray): R x 4 array of object proposals or None (for RPN)

    Returns:
        scores (ndarray): R x K array of object class scores (K includes
            background as object category 0)
        boxes (ndarray): R x (4*K) array of predicted bounding boxes
    """
    if _t:
        _t['im_preproc'].tic()
    blobs, im_scales = _get_blobs(im, boxes)
    # blob是一个ndarray, 它的轴为(batch elem, channel, height, width)
    # im_scales也是一个ndarray, 它的轴为(batch elem, im_scale)

    # When mapping from image ROIs to feature map ROIs, there's some aliasing
    # (some distinct image ROIs get mapped to the same feature ROI).
    # Here, we identify duplicate feature ROIs, so we only compute features
    # on the unique subset.
    if cfg.DEDUP_BOXES > 0 and not cfg.TEST.HAS_RPN:
        v = np.array([1, 1e3, 1e6, 1e9, 1e12])
        hashes = np.round(blobs['rois'] * cfg.DEDUP_BOXES).dot(v)
        _, index, inv_index = np.unique(hashes, return_index=True,
                                        return_inverse=True)
        blobs['rois'] = blobs['rois'][index, :]
        boxes = boxes[index, :]

    if cfg.TEST.HAS_RPN:
        im_blob = blobs['data']
        blobs['im_info'] = np.array(
            [np.hstack((im_blob.shape[2], im_blob.shape[3], im_scales[0]))],
            dtype=np.float32)


    # reshape network inputs
    net.blobs['data'].reshape(*(blobs['data'].shape))
    if cfg.TEST.HAS_RPN:
        net.blobs['im_info'].reshape(*(blobs['im_info'].shape))
    else:
        net.blobs['rois'].reshape(*(blobs['rois'].shape))

    # do forward
    net.blobs['data'].data[...] = blobs['data']
    #forward_kwargs = {'data': blobs['data'].astype(np.float32, copy=False)}
    if cfg.TEST.HAS_RPN:
        net.blobs['im_info'].data[...] = blobs['im_info']
        #forward_kwargs['im_info'] = blobs['im_info'].astype(np.float32, copy=False)
    else:
        net.blobs['rois'].data[...] = blobs['rois']
        #forward_kwargs['rois'] = blobs['rois'].astype(np.float32, copy=False)
    if _t:
        _t['im_preproc'].toc()

    if _t:
        _t['im_net'].tic()
    blobs_out = net.forward()
    if _t:
        _t['im_net'].toc()
    #blobs_out = net.forward(**forward_kwargs)

    if _t:
        _t['im_postproc'].tic()
    if cfg.TEST.HAS_RPN:
        assert len(im_scales) == 1, "Only single-image batch implemented"
        rois = net.blobs['rois'].data.copy()
        # unscale back to raw image space
        boxes = rois[:, 1:5] / im_scales[0]

    if cfg.TEST.SVM:
        # use the raw scores before softmax under the assumption they
        # were trained as linear SVMs
        scores = net.blobs['cls_score'].data
    else:
        # use softmax estimated probabilities
        scores = blobs_out['cls_prob']

    if cfg.TEST.BBOX_REG:
        # Apply bounding-box regression deltas
        box_deltas = blobs_out['bbox_pred']
        pred_boxes = bbox_transform_inv(boxes, box_deltas)
        pred_boxes = clip_boxes(pred_boxes, im.shape)
    else:
        # Simply repeat the boxes, once for each class
        pred_boxes = np.tile(boxes, (1, scores.shape[1]))

    if cfg.DEDUP_BOXES > 0 and not cfg.TEST.HAS_RPN:
        # Map scores and predictions back to the original set of boxes
        scores = scores[inv_index, :]
        pred_boxes = pred_boxes[inv_index, :]
    if _t:
        _t['im_postproc'].toc()

    return scores, pred_boxes

# def vis_detections(im, class_name, dets, thresh=0.3):
#     """Visual debugging of detections."""
#     import matplotlib.pyplot as plt
#     im = im[:, :, (2, 1, 0)]
#     for i in xrange(np.minimum(10, dets.shape[0])):
#         bbox = dets[i, :4]
#         score = dets[i, -1]
#         if score > thresh:
#             plt.cla()
#             plt.imshow(im)
#             plt.gca().add_patch(
#                 plt.Rectangle((bbox[0], bbox[1]),
#                               bbox[2] - bbox[0],
#                               bbox[3] - bbox[1], fill=False,
#                               edgecolor='g', linewidth=3)
#                 )
#             plt.title('{}  {:.3f}'.format(class_name, score))
#             plt.show()
def vis_detections(ind_color, imagename, class_name, dets, ax):
    """Visual debugging of detections."""
    edgeColor = ('red', 'green', 'blue', 'yellow')
    for i in range(dets.shape[0]):
        row = OrderedDict(
            [('ImgName', imagename), ('ClsName', class_name), ('score', dets[i, -1]),
             ('xmin', dets[i, 0]), ('ymin', dets[i, 1]), ('xmax', dets[i, 2]), ('ymax', dets[i, 3])])
        dict_list_before.append(row)
        bbox = dets[i, :4]
        score = dets[i, -1]
        ax.add_patch(
            plt.Rectangle((bbox[0], bbox[1]),
                          bbox[2] - bbox[0],
                          bbox[3] - bbox[1], fill=False,
                          edgecolor=edgeColor[ind_color], linewidth=3.5)
        )

        ax.text(bbox[0], bbox[1] - 2,
                '{:s} {:.3f}'.format(class_name[0], score),
                bbox=dict(facecolor='blue', alpha=0.5),
                fontsize=14, color='white')

    # print(dict_list_before)

def apply_nms(all_boxes, thresh):
    """Apply non-maximum suppression to all predicted boxes output by the
    test_net method.
    """
    num_classes = len(all_boxes)
    num_images = len(all_boxes[0])
    nms_boxes = [[[] for _ in xrange(num_images)]
                 for _ in xrange(num_classes)]
    for cls_ind in xrange(num_classes):
        for im_ind in xrange(num_images):
            dets = all_boxes[cls_ind][im_ind]
            if dets == []:
                continue
            # CPU NMS is much faster than GPU NMS when the number of boxes
            # is relative small (e.g., < 10k)
            # TODO(rbg): autotune NMS dispatch
            keep = nms(dets, thresh, force_cpu=True)
            if len(keep) == 0:
                continue
            nms_boxes[cls_ind][im_ind] = dets[keep, :].copy()
    return nms_boxes

def bbox_vote(dets_NMS, dets_all, thresh=0.5):
    dets_voted = np.zeros_like(dets_NMS)   # Empty matrix with the same shape and type

    _overlaps = bbox_overlaps(
			      np.ascontiguousarray(dets_NMS[:, 0:4], dtype=np.float),
			      np.ascontiguousarray(dets_all[:, 0:4], dtype=np.float))

    # for each survived box
    for i, det in enumerate(dets_NMS):
        dets_overlapped = dets_all[np.where(_overlaps[i, :] >= thresh)[0]]
        assert(len(dets_overlapped) > 0)

        boxes = dets_overlapped[:, 0:4]
        scores = dets_overlapped[:, 4]

        out_box = np.dot(scores, boxes)

        dets_voted[i][0:4] = out_box / sum(scores)        # Weighted bounding boxes
        dets_voted[i][4] = det[4]                         # Keep the original score

        # Weighted scores (if enabled)
        if cfg.TEST.BBOX_VOTE_N_WEIGHTED_SCORE > 1:
            n_agreement = cfg.TEST.BBOX_VOTE_N_WEIGHTED_SCORE
            w_empty = cfg.TEST.BBOX_VOTE_WEIGHT_EMPTY

            n_detected = len(scores)

            if n_detected >= n_agreement:
                top_scores = -np.sort(-scores)[:n_agreement]
                new_score = np.average(top_scores)
            else:
                new_score = np.average(scores) * (n_detected * 1.0 + (n_agreement - n_detected) * w_empty) / n_agreement

            dets_voted[i][4] = min(new_score, dets_voted[i][4])

    return dets_voted

def test_net(net, imdb, max_per_image=100, conf_thresh=0.001, nms_ovthresh=0.4, eval_ovthresh=0.5,
             vis=False, vis_eval=False):
    """Test a Fast R-CNN network on an image database.
    注意: 提高conf_thresh可能导致某些all_boxes[j][i]为空, 故在源码的基础上添加了一些空判断与非空判断,
    """
    num_images = len(imdb.image_index)
    # all detections are collected into:
    #    all_boxes[cls][image] = N x 5 array of detections in
    #    (x1, y1, x2, y2, score)
    all_boxes = [[[] for _ in xrange(num_images)]
                 for _ in xrange(imdb.num_classes)]

    output_dir = get_output_dir(imdb, net)

    # timers
    _t = {'im_preproc': Timer(), 'im_net': Timer(), 'im_postproc': Timer(), 'misc' : Timer()}

    if not cfg.TEST.HAS_RPN:
        roidb = imdb.roidb
        # print('--------roidb is {}'.format(roidb))
        # 调用如下：
        # roidb = imdb.roidb = imdb._roidb = imdb.roidb_handler()
        #       = imdb._roidb_handler() = imdb.selective_search_roidb()
        # 第一个括号是比较奇怪的, 是不是因为后面加的后面两个是我自己加的
        # 不用奇怪, 其实加括号是指调用(执行计算),返回结果

    ####### ----Newly added----设置图片画框后的存放路径------ ############
    if vis:
        current_time = time.strftime('%Y-%m-%d_%H-%M-%S', time.localtime(time.time()))
        test_result_savepath = os.path.join(cfg.DATA_DIR, "test_result", current_time
                                            + '_before-eval_' + "".join(imdb.classes[1:])
                                            + '_conf_' + str(conf_thresh) + "_nms_" + str(nms_ovthresh))
        if not os.path.isdir(test_result_savepath):
            os.makedirs(test_result_savepath)

    for i in xrange(num_images):
        # filter out any ground truth boxes
        if cfg.TEST.HAS_RPN:
            box_proposals = None
        else:
            # The roidb may contain ground-truth rois (for example, if the roidb
            # comes from the training or val split). We only want to evaluate
            # detection on the *non*-ground-truth rois. We select those the rois
            # that have the gt_classes field set to 0, which means there's no
            # ground truth.
            box_proposals = roidb[i]['boxes'][roidb[i]['gt_classes'] == 0]

        im_filepath = imdb.image_path_at(i)
        im = cv2.imread(im_filepath)
        scores, boxes = im_detect(net, im, _t, box_proposals)

        ####### ----Newly added----Set the Savepath------ ############
        if vis:
            im = im[:, :, (2, 1, 0)]
            # 函数返回一个figure图像和一个子图ax的array列表
            # figsize定义画布大小, 与原图保持一致
            fig, ax = plt.subplots(figsize=(im.shape[1] * 0.01, im.shape[0] * 0.01))
            ax.imshow(im)
            # 使得图片充满整个画布,从而使得输出的图片与原图一样
            plt.subplots_adjust(top=1, bottom=0, left=0, right=1, hspace=0, wspace=0)  # 使得

        hrs_Num = 0

        _t['misc'].tic()
        # skip j = 0, because it's the background class
        for j in xrange(1, imdb.num_classes):
            inds = np.where(scores[:, j] > conf_thresh)[0]
            # 添加空判断, 没有添加判断之前,程序自行中断,没有报错
            if len(inds) == 0:
                continue

            cls_scores = scores[inds, j]
            cls_boxes = boxes[inds, j*4:(j+1)*4]
            cls_dets = np.hstack((cls_boxes, cls_scores[:, np.newaxis])) \
                .astype(np.float32, copy=False)
            # print('---cls_dets after confidence filtering is {}---'.format(cls_dets))
            # 非极大值抑制
            # 极大值抑制的最终结果就是留下每一个可能的box,这些box的得分都尽可能的最大, 且这些box之间的重合度小于阈值
            keep = nms(cls_dets, nms_ovthresh)  # keep = nms(cls_dets, cfg.TEST.NMS)

            dets_NMSed = cls_dets[keep, :]
            # # bbox_vote会调整dets_NMSed, 输出最终的cls_dets, box坐标/得分可能会变
            # if cfg.TEST.BBOX_VOTE:
            #     cls_dets = bbox_vote(dets_NMSed, cls_dets)
            # else:
            #     cls_dets = dets_NMSed
            cls_dets = dets_NMSed
            # print('---cls_dets after nms is {}---'.format(cls_dets))

            if vis:
                ####### ----Newly added----先导出概率低于CONF_THRESH的图片------ ############
                if cls_dets.shape[0] == 0:
                    hrs_Num += 1
                    shutil.copy(im_filepath, test_result_savepath)
                    continue
                # vis_detections(im, imdb.classes[j], cls_dets)
                color_index = j - 1
                vis_detections(color_index, imdb.image_index[i], imdb.classes[j], cls_dets, ax)

            all_boxes[j][i] = cls_dets
            # 打印boxes
            # print('---{}--{}: {}----'.format(imdb.image_index[i], imdb.classes[j], all_boxes[j][i]))

        if vis:
            ####### ----Newly added---------- ############
            if hrs_Num == len(imdb.classes) - 1:
                plt.close()
                return

            plt.axis('off')  # 不显示坐标尺寸
            # plt.tight_layout()  # 紧凑显示图片，居中显示
            plt.draw()
            plt.savefig(os.path.join(test_result_savepath, imdb.image_index[i] + '.jpg'))
            plt.close()

        # 注意: 在内网服务器上, 439行报错
        # for j in xrange(1, imdb.num_classes)])
        # TypeError: list indices must be integers, not tuple
        # 这是因为all_boxes是空的, 下面的代码是在源码的基础修改得到的
        # Limit to max_per_image detections *over all classes*
        if max_per_image > 0:
            all_boxes_scores = []
            for j in xrange(1, imdb.num_classes):
                # 添加空判断, 没有添加判断之前,程序报错--TypeError: list indices must be integers, not tuple
                if len(all_boxes[j][i]) == 0:
                    continue
                all_boxes_scores.append(all_boxes[j][i][:, -1])

            # 添加非空判断
            if len(all_boxes_scores) != 0:
                image_scores = np.hstack(all_boxes_scores)
                if len(image_scores) > max_per_image:
                    image_thresh = np.sort(image_scores)[-max_per_image]
                    for j in xrange(1, imdb.num_classes):
                        if len(all_boxes[j][i]) == 0:
                            continue
                        keep = np.where(all_boxes[j][i][:, -1] >= image_thresh)[0]
                        all_boxes[j][i] = all_boxes[j][i][keep, :]
                        # 打印boxes
                        #print(all_boxes[j][i])

        _t['misc'].toc()

        print('im_detect: {:d}/{:d}  net {:.3f}s  preproc {:.3f}s  postproc {:.3f}s  misc {:.3f}s' \
              .format(i + 1, num_images, _t['im_net'].average_time,
                      _t['im_preproc'].average_time, _t['im_postproc'].average_time,
                      _t['misc'].average_time))

    det_file = os.path.join(output_dir, 'detections.pkl')
    with open(det_file, 'wb') as f:
        cPickle.dump(all_boxes, f, cPickle.HIGHEST_PROTOCOL)

    # 将图片名, 类别, 得分, box坐标, 保存csv
    if vis:
        d1 = DataFrame(dict_list_before)
        # print(d1)
        d1 = d1.sort_values(by=['ImgName', 'ClsName', 'score', 'xmin', 'ymin'])
        csvsavepath1 = os.path.join(cfg.DATA_DIR, 'test_result', 'before-eval_'
                                    + current_time + "_" + "".join(imdb.classes[1:])
                                    + '_conf_' + str(conf_thresh) + "_nms_" + str(nms_ovthresh) + '.csv')
        if not os.path.exists(csvsavepath1):
            d1.to_csv(csvsavepath1, index=False, header=True)
        shutil.move(csvsavepath1, test_result_savepath)

    print('Evaluating detections')
    imdb.evaluate_detections(all_boxes, conf_thresh, eval_ovthresh, output_dir, vis_eval)

    # Newly added


    # Newly added
    print('')
    print('--CONF_THRESH is {} ---'.format(conf_thresh))
    print('--NMS_OV_THRESH is {} ---'.format(nms_ovthresh))
    print('--EVAL_OV_THRESH is {} ---'.format(eval_ovthresh))
